package com.liuqi.openai.memory.chat;

import com.liuqi.openai.data.message.ChatMessage;
import com.liuqi.openai.data.message.SystemMessage;
import com.liuqi.openai.memory.ChatMemory;
import java.util.ArrayList;
import java.util.LinkedList;
import java.util.List;

import static com.liuqi.openai.util.ValidationUtil.ensureGreaterThanZero;
import static com.liuqi.openai.util.ValidationUtil.ensureNotNull;

/**
 * 此聊天记忆功能存储最大{@link #maxMessages}条消息。
 * 如果没有足够的空间容纳新消息，则丢弃最旧的消息。但不会丢弃{@link SystemMessage}消息
 *
 * <p>
 *    会始终保留{@link SystemMessage}。
 *   一次只能保存一个{@link SystemMessage}。
 *   如果添加了具有不同内容的新{@link SystemMessage}，它将替换之前的内容。
 * </p>
 *
 * <p>
 *     缓存器为 {@link LinkedList}实现, 是非线程安全的
 * </p>
 *
 * @author liuqi
 * @date 2025/7/19
 **/
public class MessageWindowChatMemory implements ChatMemory {

    /**
     * 缓存消息
     */
    private final List<ChatMessage> messages;

    /**
     * 最大存储消息数量
     */
    private final Integer maxMessages;

    public MessageWindowChatMemory(Integer maxMessages) {
        this.maxMessages = ensureGreaterThanZero(maxMessages, "maxMessages");
        this.messages = new LinkedList<>();
    }

    public MessageWindowChatMemory(Integer maxMessages, List<ChatMessage> messages) {
        this.maxMessages = ensureGreaterThanZero(maxMessages, "maxMessages");
        ensureNotNull(messages, "messages");
        this.messages = new LinkedList<>(messages);
        ensureCapacity();
    }

    @Override
    public void add(ChatMessage message) {
        if (messages.isEmpty()) {
            messages.add(message);
            return;
        }

        // SystemMessage 只能添加在第一条
        if (message instanceof SystemMessage) {
            if (messages.get(0) instanceof SystemMessage) {
                if (messages.get(0) == message) {
                    return;
                } else {
                    messages.remove(0);
                    messages.add(0, message);
                }
            }
        } else {
            messages.add(message);
            ensureCapacity();
        }
    }

    @Override
    public List<ChatMessage> messages() {
        return new ArrayList<>(messages);
    }

    @Override
    public void clear() {
        messages.clear();
    }

    /**
     * 确保容量可用
     */
    private void ensureCapacity() {
        while (messages.size() > maxMessages) {
            int messageToRemove = 0;
            if (messages.get(0) instanceof SystemMessage) {
                messageToRemove = 1;
            }
            messages.remove(messageToRemove);
        }
    }

    public static MessageWindowChatMemory withMaxMessages(int maxMessages) {
        return new MessageWindowChatMemory(maxMessages);
    }

    public static MessageWindowChatMemory withMaxMessages(int maxMessages, SystemMessage systemMessage) {
        MessageWindowChatMemory memory = new MessageWindowChatMemory(maxMessages);
        memory.add(systemMessage);
        return memory;
    }

    @Override
    public String toString() {
        return "MessageWindowChatMemory{" +
                "messages=" + messages +
                ", maxMessages=" + maxMessages +
                '}';
    }

}
